package rtmp_server

import (
	"context"
	"fmt"
	"log"
	"sync"

	"github.com/deepch/vdk/av/avutil"
	"github.com/deepch/vdk/av/pubsub"
	"github.com/deepch/vdk/codec/aacparser"
	"github.com/deepch/vdk/codec/h264parser"
	"github.com/deepch/vdk/format/rtmp"
)

type Server struct {
	server *rtmp.Server
	port   int

	mux      sync.Mutex
	channels map[string]*pubsub.Queue

	OnPush, OnPull func(ctx context.Context, path string)
}

// NewServer rtmp server 1935
func NewServer(port int) *Server {
	return &Server{
		port:     port,
		channels: map[string]*pubsub.Queue{},
		OnPull:   func(context.Context, string) {},
		OnPush:   func(context.Context, string) {},
	}
}

func (s *Server) Serve() error {
	addr := fmt.Sprintf(":%v", s.port)
	server := &rtmp.Server{
		Addr: addr,
		HandlePublish: func(conn *rtmp.Conn) {
			ctx, ctxFunc := context.WithCancel(context.Background())
			defer ctxFunc()

			s.OnPush(ctx, conn.URL.Path)
			s.handlePublish(conn)
		},
		HandlePlay: func(conn *rtmp.Conn) {
			ctx, ctxFunc := context.WithCancel(context.Background())
			defer ctxFunc()

			s.OnPull(ctx, conn.URL.Path)
			s.handlePlay(conn)
		},
	}

	return server.ListenAndServe()
}

func (s *Server) handlePublish(conn *rtmp.Conn) {
	streams, _ := conn.Streams()
	connPath := conn.URL.Path
	log.Printf("推流: %v", connPath)
	defer log.Printf("推流关闭: %v", connPath)

	ch := s.createChannel(connPath)
	defer s.deleteChannel(connPath)

	if false {
		for _, stream := range streams {
			switch stream := stream.(type) {
			case h264parser.CodecData:
				stream.SPS()
				stream.SPS()
				log.Printf("%#v", stream.Record)
			case aacparser.CodecData:
				log.Printf("%#v", stream.ConfigBytes)
			}
		}
	}

	_ = ch.WriteHeader(streams)
	_ = avutil.CopyPackets(ch, conn)
}

func (s *Server) handlePlay(conn *rtmp.Conn) {
	connPath := conn.URL.Path
	log.Printf("拉流: %v", connPath)
	defer func() {
		log.Printf("rtmp拉流关闭 %v", connPath)
	}()

	ch, ok := s.getChannel(connPath)
	if !ok {
		log.Printf("GetChannel fail. %v", connPath)
		return
	}

	cursor := ch.Latest()
	_ = avutil.CopyFile(conn, cursor)
}

func (s *Server) getChannel(connPath string) (*pubsub.Queue, bool) {
	s.mux.Lock()
	defer s.mux.Unlock()

	ch, ok := s.channels[connPath]
	return ch, ok
}

func (s *Server) createChannel(connPath string) *pubsub.Queue {
	s.mux.Lock()
	defer s.mux.Unlock()

	ch, ok := s.channels[connPath]
	if ok {
		return ch
	}

	ch = pubsub.NewQueue()
	s.channels[connPath] = ch

	return ch
}

func (s *Server) deleteChannel(connPath string) {
	s.mux.Lock()
	defer s.mux.Unlock()

	ch, ok := s.channels[connPath]
	if ok {
		_ = ch.Close()
		delete(s.channels, connPath)
	}
}

func (s *Server) GetSessionList() []string {
	s.mux.Lock()
	defer s.mux.Unlock()

	var pathList []string
	for path, _ := range s.channels {
		pathList = append(pathList, path)
	}

	return pathList
}
